{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "3fb1cef5-1a11-4129-9acf-c84f234107de",
   "metadata": {},
   "source": [
    "# LoRA PEFT"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "40fad5ea-24d6-4278-bc83-23da7f7e78d2",
   "metadata": {},
   "source": [
    "## PEFT 库 LoRA 实战 - OpenAI Whisper-Large-v3\n",
    "\n",
    "使用 LoRA 在`OpenAI Whisper-Large-v3`模型上实现`语音识别(ASR)`任务的微调训练。\n",
    "\n",
    "同时，我们还结合了`int8` 量化进一步降低训练过程资源开销，同时保证了精度几乎不受影响。\n",
    "\n",
    "主要训练流程如下：\n",
    "- 全局参数设置\n",
    "- 数据准备\n",
    "    - 下载数据集：训练、验证和评估集\n",
    "    - 预处理数据：降采样、移除不必要字段等\n",
    "    - 数据抽样（演示需要）\n",
    "    - 应用数据集处理（`Dataset.map`）\n",
    "    - 自定义语音数据处理器\n",
    "- 模型准备\n",
    "    - 加载和处理 `int8` 精度 Whisper-Large-v3 模型\n",
    "    - LoRA Adapter 参数配置\n",
    "    - 实例化 PEFT Model：`peft_model = get_peft_model(model, config)`\n",
    "- 模型训练\n",
    "    - 训练参数配置 Seq2SeqTrainingArguments\n",
    "    - 实例化训练器 Seq2SeqTrainer\n",
    "    - 训练模型\n",
    "    - 保存模型\n",
    "- 模型推理\n",
    "    - 使用 `PeftModel` 加载 LoRA 微调后 Whisper 模型\n",
    "    - 使用 `Pipeline API` 部署微调后 Whisper 实现中文语音识别任务\n",
    " \n",
    "对全量数据完成一次训练（尝试过设置epochs为3、5、10，但是训练时间较长，且容易中断），使用中文数据集，其他语言后续有时间再尝试  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "53ba1dbc-1725-42e3-899c-fd553cedd000",
   "metadata": {},
   "outputs": [],
   "source": [
    "## 全局参数设置"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "75a0e1a7-691f-4133-a204-d4d6bb3679f5",
   "metadata": {},
   "outputs": [],
   "source": [
    "model_name_or_path = \"openai/whisper-Large-v3\"\n",
    "model_dir = \"models/whisper-Large-v3-asr-int8\"\n",
    "\n",
    "language = \"Chinese (China)\"\n",
    "language_abbr = \"zh-CN\"\n",
    "task = \"transcribe\"\n",
    "dataset_name = \"mozilla-foundation/common_voice_11_0\"\n",
    "\n",
    "batch_size=64"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "53fdce52-6646-411b-a29b-12c99a5ea23d",
   "metadata": {},
   "source": [
    "## 数据准备\n",
    "\n",
    "### 下载数据集 Common Voice\n",
    "\n",
    "Common Voice 11.0 数据集包含许多不同语言的录音，总时长达数小时。\n",
    "\n",
    "本教程以中文数据为例，展示如何使用 LoRA 在 Whisper-Large-v3 上进行微调训练。\n",
    "\n",
    "首先，初始化一个DatasetDict结构，并将训练集（将训练+验证拆分为训练集）和测试集拆分好，按照中文数据集构建配置加载到内存中："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "69537eef-1b5a-4faa-85da-52a91055c787",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'client_id': '95368aab163e0387e4fd4991b4f2d8ccfbd4364bf656c860230501fd27dcedf087773e4695a6cf5de9c4f1d406d582283190d065cdfa36b0e2b060cffaca977e',\n",
       " 'path': '/root/.cache/huggingface/datasets/downloads/extracted/dcc5967c754d4c815fc005d6e297d84537028996cbcf6b34190517630cbc40b4/zh-CN_train_0/common_voice_zh-CN_33211332.mp3',\n",
       " 'audio': {'path': '/root/.cache/huggingface/datasets/downloads/extracted/dcc5967c754d4c815fc005d6e297d84537028996cbcf6b34190517630cbc40b4/zh-CN_train_0/common_voice_zh-CN_33211332.mp3',\n",
       "  'array': array([-9.09494702e-13, -2.50111043e-12, -2.04636308e-12, ...,\n",
       "          1.21667417e-05,  3.23003906e-06, -2.43063369e-07]),\n",
       "  'sampling_rate': 48000},\n",
       " 'sentence': '性喜温暖润湿气候且耐寒。',\n",
       " 'up_votes': 2,\n",
       " 'down_votes': 0,\n",
       " 'age': '',\n",
       " 'gender': '',\n",
       " 'accent': '',\n",
       " 'locale': 'zh-CN',\n",
       " 'segment': ''}"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from datasets import load_dataset, DatasetDict\n",
    "\n",
    "common_voice = DatasetDict()\n",
    "\n",
    "common_voice[\"train\"] = load_dataset(dataset_name, language_abbr, split=\"train\", trust_remote_code=True)\n",
    "common_voice[\"validation\"] = load_dataset(dataset_name, language_abbr, split=\"validation\", trust_remote_code=True)\n",
    "\n",
    "common_voice[\"train\"][0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "ca343909-d08a-4cd5-9dcf-e29b7faec34c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'client_id': '364f6dad7dcfaffc50b1cfa024509f73de36c3159ec33f57f62d9c31d37b820cf8abd26d1e760cf91c2f6e2d21f586f755cfd0365959dfbe52acdcc493910c3a',\n",
       " 'path': '/root/.cache/huggingface/datasets/downloads/extracted/f1814bfc06e59d1f106f2536b05bd9c7e529f9212c508f869745509667133820/zh-CN_dev_0/common_voice_zh-CN_18654294.mp3',\n",
       " 'audio': {'path': '/root/.cache/huggingface/datasets/downloads/extracted/f1814bfc06e59d1f106f2536b05bd9c7e529f9212c508f869745509667133820/zh-CN_dev_0/common_voice_zh-CN_18654294.mp3',\n",
       "  'array': array([ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00, ...,\n",
       "         -1.00890416e-11,  4.75411724e-12, -1.63623108e-11]),\n",
       "  'sampling_rate': 48000},\n",
       " 'sentence': '他后来跑去打她两个耳光',\n",
       " 'up_votes': 2,\n",
       " 'down_votes': 0,\n",
       " 'age': 'twenties',\n",
       " 'gender': 'male',\n",
       " 'accent': '出生地：32 江苏省',\n",
       " 'locale': 'zh-CN',\n",
       " 'segment': ''}"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "common_voice[\"validation\"][0]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "23a0380b-39f1-40ae-b679-75f6372c0640",
   "metadata": {},
   "source": [
    "## 预处理训练数据集"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "2606f0e4-56ce-4f60-bf74-fa83f57e3404",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.\n",
      "Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.\n"
     ]
    }
   ],
   "source": [
    "from transformers import AutoFeatureExtractor, AutoTokenizer, AutoProcessor\n",
    "\n",
    "# 从预训练模型加载特征提取器\n",
    "feature_extractor = AutoFeatureExtractor.from_pretrained(model_name_or_path)\n",
    "\n",
    "# 从预训练模型加载分词器，可以指定语言和任务以获得最适合特定需求的分词器配置\n",
    "tokenizer = AutoTokenizer.from_pretrained(\n",
    "    model_name_or_path, language=language, task=task)\n",
    "\n",
    "# 从预训练模型加载处理器，处理器通常结合了特征提取器和分词器，为特定任务提供一站式的数据预处理\n",
    "processor = AutoProcessor.from_pretrained(\n",
    "    model_name_or_path, language=language, task=task)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "29324ef4-7637-4698-acc6-6f4f23334838",
   "metadata": {},
   "source": [
    "#### 移除数据集中不必要的字段"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "94d1138a-ce9c-4095-a3fa-173524ae2d89",
   "metadata": {},
   "outputs": [],
   "source": [
    "common_voice = common_voice.remove_columns(\n",
    "    [\"accent\", \"age\", \"client_id\", \"down_votes\", \"gender\", \"locale\", \"path\", \"segment\", \"up_votes\"]\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "f539529c-e9a1-4669-819f-7d72cfc3a4e2",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'audio': {'path': '/root/.cache/huggingface/datasets/downloads/extracted/dcc5967c754d4c815fc005d6e297d84537028996cbcf6b34190517630cbc40b4/zh-CN_train_0/common_voice_zh-CN_33211332.mp3',\n",
       "  'array': array([-9.09494702e-13, -2.50111043e-12, -2.04636308e-12, ...,\n",
       "          1.21667417e-05,  3.23003906e-06, -2.43063369e-07]),\n",
       "  'sampling_rate': 48000},\n",
       " 'sentence': '性喜温暖润湿气候且耐寒。'}"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "common_voice[\"train\"][0]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f9564675-29be-443a-b52c-732e0a1316dd",
   "metadata": {},
   "source": [
    "#### 降采样音频数据\n",
    "\n",
    "查看`common_voice` 数据集介绍，你会发现其音频是以48kHz的采样率进行采样的.\n",
    "\n",
    "而`Whisper`模型是在16kHZ的音频输入上预训练的，因此我们需要将音频输入降采样以匹配模型预训练时使用的采样率。\n",
    "\n",
    "通过在音频列上使用`cast_column`方法，并将`sampling_rate`设置为16kHz来对音频进行降采样。\n",
    "\n",
    "下次调用时，音频输入将实时重新取样："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "6ce301a4-7d77-42e3-8a24-c5c6a8bd374a",
   "metadata": {},
   "outputs": [],
   "source": [
    "from datasets import Audio\n",
    "\n",
    "common_voice = common_voice.cast_column(\"audio\", Audio(sampling_rate=16000))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "511461da-6ad7-4f83-8dbe-a978f9c82a2d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'audio': {'path': '/root/.cache/huggingface/datasets/downloads/extracted/dcc5967c754d4c815fc005d6e297d84537028996cbcf6b34190517630cbc40b4/zh-CN_train_0/common_voice_zh-CN_33211332.mp3',\n",
       "  'array': array([ 7.27595761e-12, -5.82076609e-11, -8.73114914e-11, ...,\n",
       "         -5.96659083e-06,  2.71383469e-05,  1.29687978e-05]),\n",
       "  'sampling_rate': 16000},\n",
       " 'sentence': '性喜温暖润湿气候且耐寒。'}"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# sampling_rate 从 48KHZ 降为 16KHZ\n",
    "common_voice[\"train\"][0]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "26446056-d2c4-4710-a68f-0766e0badf31",
   "metadata": {},
   "source": [
    "### 整合以上数据处理为一个函数\n",
    "\n",
    "该数据预处理函数应该包括：\n",
    "- 通过加载音频列将音频输入重新采样为16kHZ。\n",
    "- 使用特征提取器从音频数组计算输入特征。\n",
    "- 将句子列标记化为输入标签。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "6c68f36c-5996-4c3e-a3df-a32b7fb6f1af",
   "metadata": {},
   "outputs": [],
   "source": [
    "def prepare_dataset(batch):\n",
    "    audio = batch[\"audio\"]\n",
    "    batch[\"input_features\"] = feature_extractor(audio[\"array\"], sampling_rate=audio[\"sampling_rate\"]).input_features[0]\n",
    "    batch[\"labels\"] = tokenizer(batch[\"sentence\"]).input_ids\n",
    "    return batch"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "750d920f-706b-4d63-bf3b-bb44cbfae6ce",
   "metadata": {
    "jupyter": {
     "source_hidden": true
    }
   },
   "source": [
    "### 数据抽样（课程演示DEMO）\n",
    "\n",
    "在 Whisper-Large-v3 上使用小规模数据进行演示训练，保持以下训练参数不变（batch_size=64）。\n",
    "\n",
    "使用 640 个样本训练，320个样本验证和评估，恰好使得1个 epoch 仅需10 steps 即可完成训练。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "5aaa62ce-4651-4e34-b1e8-8b11663367df",
   "metadata": {},
   "outputs": [],
   "source": [
    "small_common_voice = DatasetDict()\n",
    "\n",
    "small_common_voice[\"train\"] = common_voice[\"train\"].shuffle(seed=16).select(range(640))\n",
    "small_common_voice[\"validation\"] = common_voice[\"validation\"].shuffle(seed=16).select(range(320))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "ce14ae70-c6de-4307-93fe-2906a662f8de",
   "metadata": {
    "collapsed": true,
    "jupyter": {
     "outputs_hidden": true,
     "source_hidden": true
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DatasetDict({\n",
       "    train: Dataset({\n",
       "        features: ['audio', 'sentence'],\n",
       "        num_rows: 640\n",
       "    })\n",
       "    validation: Dataset({\n",
       "        features: ['audio', 'sentence'],\n",
       "        num_rows: 320\n",
       "    })\n",
       "})"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "small_common_voice"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "703bf15e-4577-4b12-8f42-a7c7f0de9e94",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 抽样数据处理\n",
    "# tokenized_common_voice = small_common_voice.map(prepare_dataset)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b883a56b-56b6-4bcf-8c45-e8e42ac39b07",
   "metadata": {},
   "source": [
    "## 全量训练"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "47603be3-834c-475b-a5d2-a0e9056ecd26",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DatasetDict({\n",
       "    train: Dataset({\n",
       "        features: ['audio', 'sentence'],\n",
       "        num_rows: 29056\n",
       "    })\n",
       "    validation: Dataset({\n",
       "        features: ['audio', 'sentence'],\n",
       "        num_rows: 10581\n",
       "    })\n",
       "})"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# small_common_voice = DatasetDict()\n",
    "\n",
    "# small_common_voice[\"train\"] = common_voice[\"train\"].shuffle(seed=16)\n",
    "# small_common_voice[\"validation\"] = common_voice[\"validation\"].shuffle(seed=16)\n",
    "\n",
    "common_voice"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "106114f1-35f1-4871-8025-6892632ebaef",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "05f64dc3626045f2b0a2e84d0205f27c",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map (num_proc=16):   0%|          | 0/29056 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mRuntimeError\u001b[0m                              Traceback (most recent call last)",
      "File \u001b[0;32m/usr/local/miniconda3/lib/python3.11/site-packages/datasets/arrow_dataset.py:3197\u001b[0m, in \u001b[0;36mDataset.map\u001b[0;34m(self, function, with_indices, with_rank, input_columns, batched, batch_size, drop_last_batch, remove_columns, keep_in_memory, load_from_cache_file, cache_file_name, writer_batch_size, features, disable_nullable, fn_kwargs, num_proc, suffix_template, new_fingerprint, desc)\u001b[0m\n\u001b[1;32m   3192\u001b[0m \u001b[38;5;28;01mwith\u001b[39;00m hf_tqdm(\n\u001b[1;32m   3193\u001b[0m     unit\u001b[38;5;241m=\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m examples\u001b[39m\u001b[38;5;124m\"\u001b[39m,\n\u001b[1;32m   3194\u001b[0m     total\u001b[38;5;241m=\u001b[39mpbar_total,\n\u001b[1;32m   3195\u001b[0m     desc\u001b[38;5;241m=\u001b[39m(desc \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mMap\u001b[39m\u001b[38;5;124m\"\u001b[39m) \u001b[38;5;241m+\u001b[39m \u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m (num_proc=\u001b[39m\u001b[38;5;132;01m{\u001b[39;00mnum_proc\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m)\u001b[39m\u001b[38;5;124m\"\u001b[39m,\n\u001b[1;32m   3196\u001b[0m ) \u001b[38;5;28;01mas\u001b[39;00m pbar:\n\u001b[0;32m-> 3197\u001b[0m     \u001b[38;5;28;01mfor\u001b[39;00m rank, done, content \u001b[38;5;129;01min\u001b[39;00m iflatmap_unordered(\n\u001b[1;32m   3198\u001b[0m         pool, Dataset\u001b[38;5;241m.\u001b[39m_map_single, kwargs_iterable\u001b[38;5;241m=\u001b[39mkwargs_per_job\n\u001b[1;32m   3199\u001b[0m     ):\n\u001b[1;32m   3200\u001b[0m         \u001b[38;5;28;01mif\u001b[39;00m done:\n",
      "File \u001b[0;32m/usr/local/miniconda3/lib/python3.11/site-packages/datasets/utils/py_utils.py:658\u001b[0m, in \u001b[0;36miflatmap_unordered\u001b[0;34m(pool, func, kwargs_iterable)\u001b[0m\n\u001b[1;32m    657\u001b[0m             \u001b[38;5;66;03m# One of the subprocesses has died. We should not wait forever.\u001b[39;00m\n\u001b[0;32m--> 658\u001b[0m             \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mRuntimeError\u001b[39;00m(\n\u001b[1;32m    659\u001b[0m                 \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mOne of the subprocesses has abruptly died during map operation.\u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m    660\u001b[0m                 \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mTo debug the error, disable multiprocessing.\u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m    661\u001b[0m             )\n\u001b[1;32m    662\u001b[0m \u001b[38;5;28;01mfinally\u001b[39;00m:\n",
      "\u001b[0;31mRuntimeError\u001b[0m: One of the subprocesses has abruptly died during map operation.To debug the error, disable multiprocessing.",
      "\nDuring handling of the above exception, another exception occurred:\n",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[13], line 4\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[38;5;66;03m# 完整数据训练，尝试开启 `num_proc=8` 参数多进程并行处理（如阻塞无法运行，则不使用此参数）\u001b[39;00m\n\u001b[1;32m      2\u001b[0m \u001b[38;5;66;03m# tokenized_common_voice = common_voice.map(prepare_dataset, num_proc=8)\u001b[39;00m\n\u001b[0;32m----> 4\u001b[0m tokenized_common_voice \u001b[38;5;241m=\u001b[39m common_voice\u001b[38;5;241m.\u001b[39mmap(prepare_dataset, num_proc\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m16\u001b[39m)\n",
      "File \u001b[0;32m/usr/local/miniconda3/lib/python3.11/site-packages/datasets/dataset_dict.py:868\u001b[0m, in \u001b[0;36mDatasetDict.map\u001b[0;34m(self, function, with_indices, with_rank, input_columns, batched, batch_size, drop_last_batch, remove_columns, keep_in_memory, load_from_cache_file, cache_file_names, writer_batch_size, features, disable_nullable, fn_kwargs, num_proc, desc)\u001b[0m\n\u001b[1;32m    865\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m cache_file_names \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[1;32m    866\u001b[0m     cache_file_names \u001b[38;5;241m=\u001b[39m {k: \u001b[38;5;28;01mNone\u001b[39;00m \u001b[38;5;28;01mfor\u001b[39;00m k \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mself\u001b[39m}\n\u001b[1;32m    867\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m DatasetDict(\n\u001b[0;32m--> 868\u001b[0m     {\n\u001b[1;32m    869\u001b[0m         k: dataset\u001b[38;5;241m.\u001b[39mmap(\n\u001b[1;32m    870\u001b[0m             function\u001b[38;5;241m=\u001b[39mfunction,\n\u001b[1;32m    871\u001b[0m             with_indices\u001b[38;5;241m=\u001b[39mwith_indices,\n\u001b[1;32m    872\u001b[0m             with_rank\u001b[38;5;241m=\u001b[39mwith_rank,\n\u001b[1;32m    873\u001b[0m             input_columns\u001b[38;5;241m=\u001b[39minput_columns,\n\u001b[1;32m    874\u001b[0m             batched\u001b[38;5;241m=\u001b[39mbatched,\n\u001b[1;32m    875\u001b[0m             batch_size\u001b[38;5;241m=\u001b[39mbatch_size,\n\u001b[1;32m    876\u001b[0m             drop_last_batch\u001b[38;5;241m=\u001b[39mdrop_last_batch,\n\u001b[1;32m    877\u001b[0m             remove_columns\u001b[38;5;241m=\u001b[39mremove_columns,\n\u001b[1;32m    878\u001b[0m             keep_in_memory\u001b[38;5;241m=\u001b[39mkeep_in_memory,\n\u001b[1;32m    879\u001b[0m             load_from_cache_file\u001b[38;5;241m=\u001b[39mload_from_cache_file,\n\u001b[1;32m    880\u001b[0m             cache_file_name\u001b[38;5;241m=\u001b[39mcache_file_names[k],\n\u001b[1;32m    881\u001b[0m             writer_batch_size\u001b[38;5;241m=\u001b[39mwriter_batch_size,\n\u001b[1;32m    882\u001b[0m             features\u001b[38;5;241m=\u001b[39mfeatures,\n\u001b[1;32m    883\u001b[0m             disable_nullable\u001b[38;5;241m=\u001b[39mdisable_nullable,\n\u001b[1;32m    884\u001b[0m             fn_kwargs\u001b[38;5;241m=\u001b[39mfn_kwargs,\n\u001b[1;32m    885\u001b[0m             num_proc\u001b[38;5;241m=\u001b[39mnum_proc,\n\u001b[1;32m    886\u001b[0m             desc\u001b[38;5;241m=\u001b[39mdesc,\n\u001b[1;32m    887\u001b[0m         )\n\u001b[1;32m    888\u001b[0m         \u001b[38;5;28;01mfor\u001b[39;00m k, dataset \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mitems()\n\u001b[1;32m    889\u001b[0m     }\n\u001b[1;32m    890\u001b[0m )\n",
      "File \u001b[0;32m/usr/local/miniconda3/lib/python3.11/site-packages/datasets/dataset_dict.py:869\u001b[0m, in \u001b[0;36m<dictcomp>\u001b[0;34m(.0)\u001b[0m\n\u001b[1;32m    865\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m cache_file_names \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[1;32m    866\u001b[0m     cache_file_names \u001b[38;5;241m=\u001b[39m {k: \u001b[38;5;28;01mNone\u001b[39;00m \u001b[38;5;28;01mfor\u001b[39;00m k \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mself\u001b[39m}\n\u001b[1;32m    867\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m DatasetDict(\n\u001b[1;32m    868\u001b[0m     {\n\u001b[0;32m--> 869\u001b[0m         k: dataset\u001b[38;5;241m.\u001b[39mmap(\n\u001b[1;32m    870\u001b[0m             function\u001b[38;5;241m=\u001b[39mfunction,\n\u001b[1;32m    871\u001b[0m             with_indices\u001b[38;5;241m=\u001b[39mwith_indices,\n\u001b[1;32m    872\u001b[0m             with_rank\u001b[38;5;241m=\u001b[39mwith_rank,\n\u001b[1;32m    873\u001b[0m             input_columns\u001b[38;5;241m=\u001b[39minput_columns,\n\u001b[1;32m    874\u001b[0m             batched\u001b[38;5;241m=\u001b[39mbatched,\n\u001b[1;32m    875\u001b[0m             batch_size\u001b[38;5;241m=\u001b[39mbatch_size,\n\u001b[1;32m    876\u001b[0m             drop_last_batch\u001b[38;5;241m=\u001b[39mdrop_last_batch,\n\u001b[1;32m    877\u001b[0m             remove_columns\u001b[38;5;241m=\u001b[39mremove_columns,\n\u001b[1;32m    878\u001b[0m             keep_in_memory\u001b[38;5;241m=\u001b[39mkeep_in_memory,\n\u001b[1;32m    879\u001b[0m             load_from_cache_file\u001b[38;5;241m=\u001b[39mload_from_cache_file,\n\u001b[1;32m    880\u001b[0m             cache_file_name\u001b[38;5;241m=\u001b[39mcache_file_names[k],\n\u001b[1;32m    881\u001b[0m             writer_batch_size\u001b[38;5;241m=\u001b[39mwriter_batch_size,\n\u001b[1;32m    882\u001b[0m             features\u001b[38;5;241m=\u001b[39mfeatures,\n\u001b[1;32m    883\u001b[0m             disable_nullable\u001b[38;5;241m=\u001b[39mdisable_nullable,\n\u001b[1;32m    884\u001b[0m             fn_kwargs\u001b[38;5;241m=\u001b[39mfn_kwargs,\n\u001b[1;32m    885\u001b[0m             num_proc\u001b[38;5;241m=\u001b[39mnum_proc,\n\u001b[1;32m    886\u001b[0m             desc\u001b[38;5;241m=\u001b[39mdesc,\n\u001b[1;32m    887\u001b[0m         )\n\u001b[1;32m    888\u001b[0m         \u001b[38;5;28;01mfor\u001b[39;00m k, dataset \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mitems()\n\u001b[1;32m    889\u001b[0m     }\n\u001b[1;32m    890\u001b[0m )\n",
      "File \u001b[0;32m/usr/local/miniconda3/lib/python3.11/site-packages/datasets/arrow_dataset.py:593\u001b[0m, in \u001b[0;36mtransmit_tasks.<locals>.wrapper\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m    591\u001b[0m     \u001b[38;5;28mself\u001b[39m: \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mDataset\u001b[39m\u001b[38;5;124m\"\u001b[39m \u001b[38;5;241m=\u001b[39m kwargs\u001b[38;5;241m.\u001b[39mpop(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mself\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[1;32m    592\u001b[0m \u001b[38;5;66;03m# apply actual function\u001b[39;00m\n\u001b[0;32m--> 593\u001b[0m out: Union[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mDataset\u001b[39m\u001b[38;5;124m\"\u001b[39m, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mDatasetDict\u001b[39m\u001b[38;5;124m\"\u001b[39m] \u001b[38;5;241m=\u001b[39m func(\u001b[38;5;28mself\u001b[39m, \u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)\n\u001b[1;32m    594\u001b[0m datasets: List[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mDataset\u001b[39m\u001b[38;5;124m\"\u001b[39m] \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mlist\u001b[39m(out\u001b[38;5;241m.\u001b[39mvalues()) \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28misinstance\u001b[39m(out, \u001b[38;5;28mdict\u001b[39m) \u001b[38;5;28;01melse\u001b[39;00m [out]\n\u001b[1;32m    595\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m dataset \u001b[38;5;129;01min\u001b[39;00m datasets:\n\u001b[1;32m    596\u001b[0m     \u001b[38;5;66;03m# Remove task templates if a column mapping of the template is no longer valid\u001b[39;00m\n",
      "File \u001b[0;32m/usr/local/miniconda3/lib/python3.11/site-packages/datasets/arrow_dataset.py:558\u001b[0m, in \u001b[0;36mtransmit_format.<locals>.wrapper\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m    551\u001b[0m self_format \u001b[38;5;241m=\u001b[39m {\n\u001b[1;32m    552\u001b[0m     \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mtype\u001b[39m\u001b[38;5;124m\"\u001b[39m: \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_format_type,\n\u001b[1;32m    553\u001b[0m     \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mformat_kwargs\u001b[39m\u001b[38;5;124m\"\u001b[39m: \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_format_kwargs,\n\u001b[1;32m    554\u001b[0m     \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mcolumns\u001b[39m\u001b[38;5;124m\"\u001b[39m: \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_format_columns,\n\u001b[1;32m    555\u001b[0m     \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124moutput_all_columns\u001b[39m\u001b[38;5;124m\"\u001b[39m: \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_output_all_columns,\n\u001b[1;32m    556\u001b[0m }\n\u001b[1;32m    557\u001b[0m \u001b[38;5;66;03m# apply actual function\u001b[39;00m\n\u001b[0;32m--> 558\u001b[0m out: Union[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mDataset\u001b[39m\u001b[38;5;124m\"\u001b[39m, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mDatasetDict\u001b[39m\u001b[38;5;124m\"\u001b[39m] \u001b[38;5;241m=\u001b[39m func(\u001b[38;5;28mself\u001b[39m, \u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)\n\u001b[1;32m    559\u001b[0m datasets: List[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mDataset\u001b[39m\u001b[38;5;124m\"\u001b[39m] \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mlist\u001b[39m(out\u001b[38;5;241m.\u001b[39mvalues()) \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28misinstance\u001b[39m(out, \u001b[38;5;28mdict\u001b[39m) \u001b[38;5;28;01melse\u001b[39;00m [out]\n\u001b[1;32m    560\u001b[0m \u001b[38;5;66;03m# re-apply format to the output\u001b[39;00m\n",
      "File \u001b[0;32m/usr/local/miniconda3/lib/python3.11/site-packages/datasets/arrow_dataset.py:3189\u001b[0m, in \u001b[0;36mDataset.map\u001b[0;34m(self, function, with_indices, with_rank, input_columns, batched, batch_size, drop_last_batch, remove_columns, keep_in_memory, load_from_cache_file, cache_file_name, writer_batch_size, features, disable_nullable, fn_kwargs, num_proc, suffix_template, new_fingerprint, desc)\u001b[0m\n\u001b[1;32m   3185\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mlen\u001b[39m(kwargs_per_job) \u001b[38;5;241m<\u001b[39m num_shards:\n\u001b[1;32m   3186\u001b[0m     logger\u001b[38;5;241m.\u001b[39minfo(\n\u001b[1;32m   3187\u001b[0m         \u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mReprocessing \u001b[39m\u001b[38;5;132;01m{\u001b[39;00m\u001b[38;5;28mlen\u001b[39m(kwargs_per_job)\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m/\u001b[39m\u001b[38;5;132;01m{\u001b[39;00mnum_shards\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m shards because some of them were missing from the cache.\u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m   3188\u001b[0m     )\n\u001b[0;32m-> 3189\u001b[0m \u001b[38;5;28;01mwith\u001b[39;00m Pool(\u001b[38;5;28mlen\u001b[39m(kwargs_per_job)) \u001b[38;5;28;01mas\u001b[39;00m pool:\n\u001b[1;32m   3190\u001b[0m     os\u001b[38;5;241m.\u001b[39menviron \u001b[38;5;241m=\u001b[39m prev_env\n\u001b[1;32m   3191\u001b[0m     logger\u001b[38;5;241m.\u001b[39minfo(\u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mSpawning \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mnum_proc\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m processes\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n",
      "File \u001b[0;32m/usr/local/miniconda3/lib/python3.11/site-packages/multiprocess/pool.py:739\u001b[0m, in \u001b[0;36mPool.__exit__\u001b[0;34m(self, exc_type, exc_val, exc_tb)\u001b[0m\n\u001b[1;32m    738\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m__exit__\u001b[39m(\u001b[38;5;28mself\u001b[39m, exc_type, exc_val, exc_tb):\n\u001b[0;32m--> 739\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mterminate()\n",
      "File \u001b[0;32m/usr/local/miniconda3/lib/python3.11/site-packages/multiprocess/pool.py:657\u001b[0m, in \u001b[0;36mPool.terminate\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m    655\u001b[0m util\u001b[38;5;241m.\u001b[39mdebug(\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mterminating pool\u001b[39m\u001b[38;5;124m'\u001b[39m)\n\u001b[1;32m    656\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_state \u001b[38;5;241m=\u001b[39m TERMINATE\n\u001b[0;32m--> 657\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_terminate()\n",
      "File \u001b[0;32m/usr/local/miniconda3/lib/python3.11/site-packages/multiprocess/util.py:224\u001b[0m, in \u001b[0;36mFinalize.__call__\u001b[0;34m(self, wr, _finalizer_registry, sub_debug, getpid)\u001b[0m\n\u001b[1;32m    221\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m    222\u001b[0m     sub_debug(\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mfinalizer calling \u001b[39m\u001b[38;5;132;01m%s\u001b[39;00m\u001b[38;5;124m with args \u001b[39m\u001b[38;5;132;01m%s\u001b[39;00m\u001b[38;5;124m and kwargs \u001b[39m\u001b[38;5;132;01m%s\u001b[39;00m\u001b[38;5;124m'\u001b[39m,\n\u001b[1;32m    223\u001b[0m               \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_callback, \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_args, \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_kwargs)\n\u001b[0;32m--> 224\u001b[0m     res \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_callback(\u001b[38;5;241m*\u001b[39m\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_args, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39m\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_kwargs)\n\u001b[1;32m    225\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_weakref \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_callback \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_args \u001b[38;5;241m=\u001b[39m \\\n\u001b[1;32m    226\u001b[0m                 \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_kwargs \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_key \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m\n\u001b[1;32m    227\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m res\n",
      "File \u001b[0;32m/usr/local/miniconda3/lib/python3.11/site-packages/multiprocess/pool.py:732\u001b[0m, in \u001b[0;36mPool._terminate_pool\u001b[0;34m(cls, taskqueue, inqueue, outqueue, pool, change_notifier, worker_handler, task_handler, result_handler, cache)\u001b[0m\n\u001b[1;32m    729\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m p\u001b[38;5;241m.\u001b[39mis_alive():\n\u001b[1;32m    730\u001b[0m     \u001b[38;5;66;03m# worker has not yet exited\u001b[39;00m\n\u001b[1;32m    731\u001b[0m     util\u001b[38;5;241m.\u001b[39mdebug(\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mcleaning up worker \u001b[39m\u001b[38;5;132;01m%d\u001b[39;00m\u001b[38;5;124m'\u001b[39m \u001b[38;5;241m%\u001b[39m p\u001b[38;5;241m.\u001b[39mpid)\n\u001b[0;32m--> 732\u001b[0m     p\u001b[38;5;241m.\u001b[39mjoin()\n",
      "File \u001b[0;32m/usr/local/miniconda3/lib/python3.11/site-packages/multiprocess/process.py:149\u001b[0m, in \u001b[0;36mBaseProcess.join\u001b[0;34m(self, timeout)\u001b[0m\n\u001b[1;32m    147\u001b[0m \u001b[38;5;28;01massert\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_parent_pid \u001b[38;5;241m==\u001b[39m os\u001b[38;5;241m.\u001b[39mgetpid(), \u001b[38;5;124m'\u001b[39m\u001b[38;5;124mcan only join a child process\u001b[39m\u001b[38;5;124m'\u001b[39m\n\u001b[1;32m    148\u001b[0m \u001b[38;5;28;01massert\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_popen \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m, \u001b[38;5;124m'\u001b[39m\u001b[38;5;124mcan only join a started process\u001b[39m\u001b[38;5;124m'\u001b[39m\n\u001b[0;32m--> 149\u001b[0m res \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_popen\u001b[38;5;241m.\u001b[39mwait(timeout)\n\u001b[1;32m    150\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m res \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[1;32m    151\u001b[0m     _children\u001b[38;5;241m.\u001b[39mdiscard(\u001b[38;5;28mself\u001b[39m)\n",
      "File \u001b[0;32m/usr/local/miniconda3/lib/python3.11/site-packages/multiprocess/popen_fork.py:43\u001b[0m, in \u001b[0;36mPopen.wait\u001b[0;34m(self, timeout)\u001b[0m\n\u001b[1;32m     41\u001b[0m             \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m\n\u001b[1;32m     42\u001b[0m     \u001b[38;5;66;03m# This shouldn't block if wait() returned successfully.\u001b[39;00m\n\u001b[0;32m---> 43\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mpoll(os\u001b[38;5;241m.\u001b[39mWNOHANG \u001b[38;5;28;01mif\u001b[39;00m timeout \u001b[38;5;241m==\u001b[39m \u001b[38;5;241m0.0\u001b[39m \u001b[38;5;28;01melse\u001b[39;00m \u001b[38;5;241m0\u001b[39m)\n\u001b[1;32m     44\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mreturncode\n",
      "File \u001b[0;32m/usr/local/miniconda3/lib/python3.11/site-packages/multiprocess/popen_fork.py:27\u001b[0m, in \u001b[0;36mPopen.poll\u001b[0;34m(self, flag)\u001b[0m\n\u001b[1;32m     25\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mreturncode \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[1;32m     26\u001b[0m     \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m---> 27\u001b[0m         pid, sts \u001b[38;5;241m=\u001b[39m os\u001b[38;5;241m.\u001b[39mwaitpid(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mpid, flag)\n\u001b[1;32m     28\u001b[0m     \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mOSError\u001b[39;00m:\n\u001b[1;32m     29\u001b[0m         \u001b[38;5;66;03m# Child process not yet created. See #1731717\u001b[39;00m\n\u001b[1;32m     30\u001b[0m         \u001b[38;5;66;03m# e.errno == errno.ECHILD == 10\u001b[39;00m\n\u001b[1;32m     31\u001b[0m         \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m\n",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "# 完整数据训练，尝试开启 `num_proc=8` 参数多进程并行处理（如阻塞无法运行，则不使用此参数）\n",
    "# tokenized_common_voice = common_voice.map(prepare_dataset, num_proc=8)\n",
    "\n",
    "tokenized_common_voice = common_voice.map(prepare_dataset, num_proc=16)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "8d528bd8-0524-45c2-8761-c6b57db66da4",
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenized_common_voice = common_voice.map(prepare_dataset) "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f9a4969a-4bf5-4541-80c7-7f3158e63085",
   "metadata": {},
   "source": [
    "## 自定义语音数据整理器\n",
    "\n",
    "定义了一个针对语音到文本（Seq2Seq）模型的自定义数据整理器类，特别适用于输入为语音特征、输出为文本序列的数据集。\n",
    "\n",
    "\n",
    "这个整理器（`DataCollatorSpeechSeq2SeqWithPadding`）旨在将数据点批量打包，将每个批次中的`attention_mask`填充到最大长度，以保持批处理中张量形状的一致性，并用`-100`替换填充值，以便在损失函数中被忽略。这对于神经网络的高效训练至关重要。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "3abcf91a-8e4f-4345-b4a5-26075d821f37",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "\n",
    "from dataclasses import dataclass\n",
    "from typing import Any, Dict, List, Union\n",
    "\n",
    "# 定义一个针对语音到文本任务的数据整理器类\n",
    "@dataclass\n",
    "class DataCollatorSpeechSeq2SeqWithPadding:\n",
    "    processor: Any  # 处理器结合了特征提取器和分词器\n",
    "\n",
    "    # 整理器函数，将特征列表处理成一个批次\n",
    "    def __call__(self, features: List[Dict[str, Union[List[int], torch.Tensor]]]) -> Dict[str, torch.Tensor]:\n",
    "        # 从特征列表中提取输入特征，并填充以使它们具有相同的形状\n",
    "        input_features = [{\"input_features\": feature[\"input_features\"]} for feature in features]\n",
    "        batch = self.processor.feature_extractor.pad(input_features, return_tensors=\"pt\")\n",
    "\n",
    "        # 从特征列表中提取标签特征（文本令牌），并进行填充\n",
    "        label_features = [{\"input_ids\": feature[\"labels\"]} for feature in features]\n",
    "        labels_batch = self.processor.tokenizer.pad(label_features, return_tensors=\"pt\")\n",
    "\n",
    "        # 使用-100替换标签中的填充区域，-100通常用于在损失计算中忽略填充令牌\n",
    "        labels = labels_batch[\"input_ids\"].masked_fill(labels_batch.attention_mask.ne(1), -100)\n",
    "\n",
    "        # 如果批次中的所有序列都以句子开始令牌开头，则移除它\n",
    "        if (labels[:, 0] == self.processor.tokenizer.bos_token_id).all().cpu().item():\n",
    "            labels = labels[:, 1:]\n",
    "\n",
    "        # 将处理过的标签添加到批次中\n",
    "        batch[\"labels\"] = labels\n",
    "\n",
    "        return batch  # 返回最终的批次，准备好进行训练或评估"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "d86d3e64-2832-4fad-a623-74a04d03aed6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 用给定的处理器实例化数据整理器\n",
    "data_collator = DataCollatorSpeechSeq2SeqWithPadding(processor=processor)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "466debe3-9f24-439d-822f-40ef1ecfa5e2",
   "metadata": {},
   "source": [
    "## 模型准备\n",
    "\n",
    "### 加载预训练模型（int8 精度）\n",
    "\n",
    "使用 `int8 ` 精度加载预训练模型，进一步降低显存需求。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "b595fd0e-a02d-4f1c-ad21-3a347006abf3",
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import AutoModelForSpeechSeq2Seq\n",
    "\n",
    "model = AutoModelForSpeechSeq2Seq.from_pretrained(model_name_or_path, load_in_8bit=True, device_map=\"auto\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "b7112ed1-b31f-4b07-ab11-eefff7f5dc8b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 设置模型配置中的forced_decoder_ids属性为None\n",
    "model.config.forced_decoder_ids = None  # 这通常用于指定在解码（生成文本）过程中必须使用的特定token的ID，设置为None表示没有这样的强制要求\n",
    "\n",
    "# 设置模型配置中的suppress_tokens列表为空\n",
    "model.config.suppress_tokens = []  # 这用于指定在生成过程中应被抑制（不生成）的token的列表，设置为空列表表示没有要抑制的token"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "518f6e67-8532-4ea3-8b1a-b2a79855c080",
   "metadata": {},
   "source": [
    "### PEFT 微调前的模型处理\n",
    "\n",
    "在使用 `peft` 训练 int8 模型之前，需要进行一些预处理：\n",
    "- 将所有非 `int8` 精度模块转换为全精度（`fp32`）以保证稳定性\n",
    "- 为输入嵌入层添加一个 `forward_hook`，以启用输入隐藏状态的梯度计算\n",
    "- 启用梯度检查点以实现更高效的内存训练\n",
    "\n",
    "使用 `peft` 库预定义的工具函数 `prepare_model_for_int8_training`，便可自动完成以上模型处理工作。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "53820452-d56d-43ab-877d-085517a33ac7",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/usr/local/miniconda3/lib/python3.11/site-packages/peft/utils/other.py:143: FutureWarning: prepare_model_for_int8_training is deprecated and will be removed in a future version. Use prepare_model_for_kbit_training instead.\n",
      "  warnings.warn(\n"
     ]
    }
   ],
   "source": [
    "from peft import prepare_model_for_int8_training\n",
    "\n",
    "model = prepare_model_for_int8_training(model)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5b6ae3cf-b48e-4bd2-abe8-294407afa562",
   "metadata": {},
   "source": [
    "### LoRA Adapter 配置\n",
    "\n",
    "在 `peft` 中使用`LoRA`非常简捷，借助 `PeftModel`抽象，我们可以快速使用低秩适配器（LoRA）到任意模型。\n",
    "\n",
    "通过使用 `peft` 中的 `get_peft_model` 工具函数来实现。\n",
    "\n",
    "#### 关于 LoRA 超参数的说明：\n",
    "```\n",
    "MatMul(B,A) * Scaling\n",
    "Scaling = LoRA_Alpha / Rank\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "3cef7c5f-86e5-4bbf-bd9a-b6c4f7bc11b2",
   "metadata": {},
   "outputs": [],
   "source": [
    "from peft import LoraConfig, PeftModel, LoraModel, LoraConfig, get_peft_model\n",
    "\n",
    "# 创建一个LoraConfig对象，用于设置LoRA（Low-Rank Adaptation）的配置参数\n",
    "config = LoraConfig(\n",
    "    r=8,  # LoRA的秩，影响LoRA矩阵的大小\n",
    "    lora_alpha=64,  # LoRA适应的比例因子\n",
    "    # 指定将LoRA应用到的模型模块，通常是attention和全连接层的投影。\n",
    "    target_modules=[\"q_proj\", \"v_proj\"],\n",
    "    lora_dropout=0.05,  # 在LoRA模块中使用的dropout率\n",
    "    bias=\"none\",  # 设置bias的使用方式，这里没有使用bias\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "deb5b305-a6eb-40a2-8ac9-ba4f39f1b930",
   "metadata": {},
   "source": [
    "### 使用get_peft_model函数和给定的配置来获取一个PEFT模型"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "2045782f-69da-4164-a469-bf6128b977c9",
   "metadata": {},
   "outputs": [],
   "source": [
    "peft_model = get_peft_model(model, config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "de4d27bb-36a4-47cb-af19-28a7fdccf6aa",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "trainable params: 3,932,160 || all params: 1,547,422,720 || trainable%: 0.2541102666503436\n"
     ]
    }
   ],
   "source": [
    "# 打印 LoRA 微调训练的模型参数\n",
    "peft_model.print_trainable_parameters()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7939d571-9077-4672-93fa-1bc6b18c86a9",
   "metadata": {},
   "source": [
    "## 模型训练\n",
    "\n",
    "#### Seq2SeqTrainingArguments 训练参数\n",
    "\n",
    "**关于设置训练步数和评估步数**\n",
    "\n",
    "基于 epochs 设置：\n",
    "\n",
    "```python\n",
    "    num_train_epochs=3,  # 训练的总轮数\n",
    "    evaluation_strategy=\"epoch\",  # 设置评估策略，这里是在每个epoch结束时进行评估\n",
    "    warmup_steps=50,  # 在训练初期增加学习率的步数，有助于稳定训练\n",
    "```\n",
    "\n",
    "基于 steps 设置：\n",
    "\n",
    "```python\n",
    "    evaluation_strategy=\"steps\", \n",
    "    eval_steps=25,\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "b911f31f-26d9-46dc-943c-8237fa1a2697",
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import Seq2SeqTrainingArguments\n",
    "\n",
    "# 设置序列到序列模型训练的参数\n",
    "training_args = Seq2SeqTrainingArguments(\n",
    "    output_dir=model_dir,  # 指定模型输出和保存的目录\n",
    "    per_device_train_batch_size=batch_size,  # 每个设备上的训练批量大小\n",
    "    learning_rate=1e-3,  # 学习率\n",
    "    num_train_epochs=1,  # 训练的总轮数\n",
    "    evaluation_strategy=\"epoch\",  # 设置评估策略，这里是在每个epoch结束时进行评估\n",
    "    warmup_steps=50,  # 在训练初期增加学习率的步数，有助于稳定训练\n",
    "    fp16=True,  # 启用混合精度训练，可以提高训练速度，同时减少内存使用\n",
    "    per_device_eval_batch_size=batch_size,  # 每个设备上的评估批量大小\n",
    "    generation_max_length=128,  # 生成任务的最大长度\n",
    "    logging_steps=10,  # 指定日志记录的步骤，用于跟踪训练进度\n",
    "    remove_unused_columns=False,  # 是否删除不使用的列，以减少数据处理开销\n",
    "    label_names=[\"labels\"],  # 指定标签列的名称，用于训练过程中\n",
    "    # evaluation_strategy=\"steps\",\n",
    "    # eval_steps=25,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "96ecab72-277a-4dcf-adb1-b027651e9176",
   "metadata": {},
   "source": [
    "### 实例化 Seq2SeqTrainer 训练器"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "62eb7ec2-8aac-4ccd-8259-091080f2e57c",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Detected kernel version 4.18.0, which is below the recommended minimum of 5.5.0; this can cause the process to hang. It is recommended to upgrade the kernel to the minimum version or higher.\n"
     ]
    }
   ],
   "source": [
    "from transformers import Seq2SeqTrainer\n",
    "\n",
    "trainer = Seq2SeqTrainer(\n",
    "    args=training_args,\n",
    "    model=peft_model,\n",
    "    train_dataset=tokenized_common_voice[\"train\"],\n",
    "    eval_dataset=tokenized_common_voice[\"validation\"],\n",
    "    data_collator=data_collator,\n",
    "    tokenizer=processor.feature_extractor,\n",
    ")\n",
    "peft_model.config.use_cache = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "9f544d55-49f4-4799-8658-be95dfd805f4",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/usr/local/miniconda3/lib/python3.11/site-packages/torch/utils/checkpoint.py:460: UserWarning: torch.utils.checkpoint: please pass in use_reentrant=True or use_reentrant=False explicitly. The default value of use_reentrant will be updated to be False in the future. To maintain current behavior, pass use_reentrant=True. It is recommended that you use use_reentrant=False. Refer to docs for more details on the differences between the two variants.\n",
      "  warnings.warn(\n",
      "/usr/local/miniconda3/lib/python3.11/site-packages/torch/utils/checkpoint.py:90: UserWarning: None of the inputs have requires_grad=True. Gradients will be None\n",
      "  warnings.warn(\n",
      "/usr/local/miniconda3/lib/python3.11/site-packages/bitsandbytes/autograd/_functions.py:322: UserWarning: MatMul8bitLt: inputs will be cast from torch.float32 to float16 during quantization\n",
      "  warnings.warn(f\"MatMul8bitLt: inputs will be cast from {A.dtype} to float16 during quantization\")\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "\n",
       "    <div>\n",
       "      \n",
       "      <progress value='454' max='454' style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
       "      [454/454 2:27:41, Epoch 1/1]\n",
       "    </div>\n",
       "    <table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       " <tr style=\"text-align: left;\">\n",
       "      <th>Epoch</th>\n",
       "      <th>Training Loss</th>\n",
       "      <th>Validation Loss</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <td>1</td>\n",
       "      <td>0.336400</td>\n",
       "      <td>0.401340</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table><p>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "TrainOutput(global_step=454, training_loss=0.4961100036877367, metrics={'train_runtime': 8888.8765, 'train_samples_per_second': 3.269, 'train_steps_per_second': 0.051, 'total_flos': 9.89809493409792e+19, 'train_loss': 0.4961100036877367, 'epoch': 1.0})"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "trainer.train()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7a021eea-afa8-44e3-95ac-ff148e5aa393",
   "metadata": {},
   "source": [
    "### 保存 LoRA 模型(Adapter)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "0a750fab-429e-4159-85e3-185e515f9c4b",
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer.save_model(model_dir)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4bd6322c-8a16-4bf7-a667-9db317fe0f02",
   "metadata": {},
   "source": [
    "## 模型推理\n",
    "\n",
    "**再次加载模型会额外占用显存，如果显存已经达到上限，建议重启 Notebook 后再进行以下操作**\n",
    "\n",
    "### 使用 `PeftModel` 加载 LoRA 微调后 Whisper 模型\n",
    "\n",
    "使用 `PeftConfig` 加载 LoRA Adapter 配置参数，使用 `PeftModel` 加载微调后 Whisper 模型"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "9ff663df-aa28-4c24-8e7f-4d50cc929307",
   "metadata": {},
   "outputs": [],
   "source": [
    "model_dir = \"models/whisper-Large-v3-asr-int8\"\n",
    "\n",
    "language = \"Chinese (China)\"\n",
    "language_abbr = \"zh-CN\"\n",
    "language_decode = \"chinese\"\n",
    "task = \"transcribe\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "a7052340-d675-4492-97f3-f1b51a97ecb1",
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import AutoModelForSpeechSeq2Seq, AutoTokenizer, AutoProcessor\n",
    "from peft import PeftConfig, PeftModel\n",
    "\n",
    "peft_config = PeftConfig.from_pretrained(model_dir)\n",
    "\n",
    "base_model = AutoModelForSpeechSeq2Seq.from_pretrained(\n",
    "    peft_config.base_model_name_or_path, load_in_8bit=True, device_map=\"auto\"\n",
    ")\n",
    "\n",
    "peft_model = PeftModel.from_pretrained(base_model, model_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "3c08cd3b-6a42-4bea-b38c-e7a49847fe09",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.\n",
      "Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.\n"
     ]
    }
   ],
   "source": [
    "tokenizer = AutoTokenizer.from_pretrained(peft_config.base_model_name_or_path, language=language, task=task)\n",
    "processor = AutoProcessor.from_pretrained(peft_config.base_model_name_or_path, language=language, task=task)\n",
    "feature_extractor = processor.feature_extractor"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f79a51a5-f6df-4251-8b78-e73df789a4e5",
   "metadata": {},
   "source": [
    "### 使用 Pipeline API 部署微调后 Whisper 实现中文语音识别任务"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "71855a90-c95b-4af8-a834-e509159b03b3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 驯养的沙特瞪羚的基因分析显示它们有可能是不同的物种或混种。\n",
    "test_audio = \"data/audio/zh_test.mp3\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "496fa062-0adb-4aed-9520-307030f2e044",
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import AutomaticSpeechRecognitionPipeline\n",
    "\n",
    "pipeline = AutomaticSpeechRecognitionPipeline(model=peft_model, tokenizer=tokenizer, feature_extractor=feature_extractor)\n",
    "\n",
    "forced_decoder_ids = processor.get_decoder_prompt_ids(language=language_decode, task=task)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "f847685f-f074-487a-b428-b4d7dbe18768",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/usr/local/miniconda3/lib/python3.11/site-packages/bitsandbytes/autograd/_functions.py:322: UserWarning: MatMul8bitLt: inputs will be cast from torch.float32 to float16 during quantization\n",
      "  warnings.warn(f\"MatMul8bitLt: inputs will be cast from {A.dtype} to float16 during quantization\")\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "\n",
    "with torch.cuda.amp.autocast():\n",
    "    text = pipeline(test_audio, max_new_tokens=255)[\"text\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "e7a7137e-58bc-4a9d-9dd5-c632c509de9a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "' projects熏养的沙特登灵的基因分析显示，它们有可能是不同的物种或混种。'"
      ]
     },
     "execution_count": 37,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "text"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0fb2a728-4b45-42ed-852a-76d64cd9d2a3",
   "metadata": {},
   "source": [
    "## 测试语音集\n",
    "\n",
    "english_mp4 = \"https://huggingface.co/datasets/Narsil/asr_dummy/resolve/main/mlk.flac\"\n",
    "text: \"I have a dream that one day this nation will rise up and live out the true meaning of its creed.\"\n",
    "\n",
    "french_mp4 = \"https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/asr_example_fr.wav\"\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
